package com.lin.project.hackathon.examples;

import com.lin.project.hackathon.ai.*;

import java.io.IOException;
import java.util.Arrays;
import java.util.List;


public class HelloWorld {
    public static void main(String... args) {
        String apiKey = "sk-kHrI842xCGau47uRW3IRnXK8G8ylLRGZf22ahsFHSPXnchv2";
        if (apiKey == null) {
            System.out.println("Please set MOONSHOT_API_KEY env");
            return;
        }
        Client client = new Client(apiKey);
        try {
            ModelsList models = client.ListModels();
            for (Model model : models.getData()) {
                System.out.println(model.getId());
            }
        } catch (IOException e) {
            e.printStackTrace();
        }

        final List<ChatCompletionMessage> messages = Arrays.asList(
                new ChatCompletionMessage(ChatMessageRole.SYSTEM.value(),
                        "你是一名高级软件工程师，擅长各种算法。"),
                new ChatCompletionMessage(ChatMessageRole.USER.value(),
                        "你好，请帮我写一个快速排序，用Java语言，并解释")
        );

//        try {
//            ChatCompletionResponse response = client.ChatCompletion(new ChatCompletionRequest(
//                    "moonshot-v1-8k",
//                    messages,
//                    50,
//                    0.3f,
//                    1
//            ));
//            for (ChatCompletionChoice choice : response.getChoices()) {
//                System.out.println(choice.getMessage().getContent());
//            }
//        } catch (IOException e) {
//            e.printStackTrace();
//        }

        StringBuilder stringBuilder = new StringBuilder();
        try {
            client.ChatCompletionStream(new ChatCompletionRequest(
                    "moonshot-v1-8k",
                    messages,
                    4096,
                    0.3f,
                    1
            )).subscribe(
                    streamResponse -> {
                        if (streamResponse.getChoices().isEmpty()) {
                            return;
                        }
                        for (ChatCompletionStreamChoice choice : streamResponse.getChoices()) {
                            String finishReason = choice.getFinishReason();
                            if (finishReason != null) {
                                System.out.println("finish reason: " + finishReason);
                                continue;
                            }
                            System.out.println(choice.getDelta().getContent());
                            stringBuilder.append(choice.getDelta().getContent());
                        }
                    },
                    error -> {
                        error.printStackTrace();
                    },
                    () -> {
                        System.out.println(stringBuilder.toString());
                        System.out.println("complete");
                    }
            );
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}
